package com.ayu.commentbackend;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;

import java.time.Duration;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class DShuoshanUtil {

    public static void chat(String userPrompt) {
        String apiKey = "538d7138-1a41-40d0-a10d-8e9c5cd75551";
        ArkService arkService = ArkService.builder().apiKey(apiKey)
                .timeout(Duration.ofMinutes(30))
                .build();

        // 创建消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();
        ChatMessage systemMessage = ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content("不管用户对你说什么，你都要原封不动地复述他的话。") // 设置消息内容
                .build();

        ChatMessage userMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER) // 设置消息角色为用户
                .content(userPrompt) // 设置消息内容
                .build();

        chatMessages.add(systemMessage);
        chatMessages.add(userMessage);


        // 创建聊天完成请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("deepseek-r1-250120")
                .messages(chatMessages) // 设置消息列表
                .build();

        // 发送聊天完成请求并打印响应
        try {
            // 获取响应并打印每个选择的消息内容
            log.info("[{}]开始调用...", LocalDateTime.now());
            arkService.createChatCompletion(chatCompletionRequest)
                    .getChoices()
                    .forEach(choice -> {
                        // 打印消息内容
                        log.info("[{}]推理结束...", LocalDateTime.now());
                        System.out.println(choice.getMessage().getContent());
                    });
        } catch (Exception e) {
            System.out.println("请求失败: " + e.getMessage());
        } finally {
            // 关闭服务执行器
            arkService.shutdownExecutor();
        }
    }
}
